import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score
from sklearn.metrics import roc_auc_score
import joblib
import transform_train


def tree_train():
    x_train, y_train = transform_train.train_do()
    x_train_tree = x_train.iloc[:, 0:20]

    model_tree = RandomForestClassifier(random_state=22)
    param_grid = {
        'n_estimators': [10,20,50,100],
        'max_depth': [3, 5, 7],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4],
        'max_features': ['sqrt', 'log2'],

    }
    grid_search = GridSearchCV(estimator=model_tree, param_grid=param_grid, cv=5, n_jobs=-1)
    grid_search.fit(x_train_tree, y_train)

    joblib.dump(grid_search.best_estimator_, '../model/tree.pkl')


if __name__ == '__main__':
    tree_train()
